import os

from zhipuai import ZhipuAI
from dotenv import load_dotenv, find_dotenv

from searchByMilvus import query_milvus_by_text

# 加载.env文件
load_dotenv()
# 从环境变量中读取API密钥
api_key = os.getenv('ZHIPU_API_KEY')
# _ = load_dotenv(find_dotenv())

client = ZhipuAI(api_key=api_key)

prompt_template = """
你是一个问答机器人。
你的任务是根据下述给定的已知信息回答用户问题。

已知信息:
{context}

用户问：
{query}

如果已知信息不包含用户问题的答案，或者已知信息不足以回答用户的问题，请直接回复"我无法回答您的问题"。
请不要输出已知信息中不包含的信息或答案。
请用中文回答用户问题。
"""


def build_prompt(prompt_template, **kwargs):
    '''将 Prompt 模板赋值'''
    inputs = {}
    for k, v in kwargs.items():
        if isinstance(v, list) and all(isinstance(elem, str) for elem in v):
            val = '\n\n'.join(v)
        else:
            val = v
        inputs[k] = val
    return prompt_template.format(**inputs)


def get_completion(prompt, model="glm-4"):
    '''封装 openai 接口'''
    messages = [{"role": "user", "content": prompt}]
    response = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=0,  # 模型输出的随机性，0 表示随机性最小
    )
    return response.choices[0].message.content


user_query = "heelo"

if __name__ == '__main__':
    # 1. 检索
    search_results = query_milvus_by_text(user_query)

    # # 2. 构建 Prompt
    # prompt = build_prompt(prompt_template, context=search_results, query=user_query)
    # print("===Prompt===")
    # print(prompt)

    # response = get_completion(prompt)

    # print("===回复===")
    # print(response)
